#ifndef __COMMON_NNIE_H__
#define __COMMON_NNIE_H__

#include "hi_types.h"
#include "hi_common.h"
#include "hi_nnie.h"
#include "hi_comm_video.h"

#define SVP_NNIE_EACH_SEG_STEP_ADDR_NUM 2

typedef struct
{
    HI_U32 u32OriImHeight;
    HI_U32 u32OriImWidth;
    HI_U32 u32BboxNumEachGrid;
    HI_U32 u32ClassNum;
    HI_U32 au32GridNumHeight[3];
    HI_U32 au32GridNumWidth[3];
    HI_U32 u32NmsThresh;
    HI_U32 u32ConfThresh;
    HI_U32 u32MaxRoiNum;
    HI_FLOAT af32Bias[3][6];
    SVP_MEM_INFO_S stGetResultTmpBuf;
    SVP_DST_BLOB_S stClassRoiNum;
    SVP_DST_BLOB_S stDstRoi;
    SVP_DST_BLOB_S stDstScore;
} SVP_YOLO_PARAM_S;

typedef struct
{
    SVP_SRC_BLOB_S astSrc[SVP_NNIE_MAX_INPUT_NUM];
    SVP_DST_BLOB_S astDst[SVP_NNIE_MAX_OUTPUT_NUM];
} SVP_NNIE_SEG_DATA_S;

typedef struct
{
    HI_U32 u32MaxInputNum;
    HI_U32 u32MaxRoiNum;

    SVP_MEM_INFO_S stModelBuf;
    SVP_NNIE_MODEL_S stModel;

    SVP_NNIE_SEG_DATA_S astSegData[SVP_NNIE_MAX_NET_SEG_NUM];

    SVP_NNIE_FORWARD_CTRL_S astForwardCtrl[SVP_NNIE_MAX_NET_SEG_NUM];
    SVP_NNIE_FORWARD_WITHBBOX_CTRL_S astForwardWithBboxCtrl[SVP_NNIE_MAX_NET_SEG_NUM];

    HI_U32 au32TaskBufSize[SVP_NNIE_MAX_NET_SEG_NUM];

    SVP_MEM_INFO_S stTaskBuf;
    SVP_MEM_INFO_S stTmpBuf;
    SVP_MEM_INFO_S stStepBuf;

    HI_U64 au64StepVirAddr[SVP_NNIE_EACH_SEG_STEP_ADDR_NUM * SVP_NNIE_MAX_NET_SEG_NUM];
} SVP_NNIE_PARAM_S;

HI_S32 SVP_NNIE_Init(HI_CHAR *pcModelName, SVP_NNIE_PARAM_S *pstNnieParam);

HI_VOID SVP_NNIE_DeInit(SVP_NNIE_PARAM_S *pstNnieParam);

HI_S32 SVP_NNIE_Forward(SVP_NNIE_PARAM_S *pstNnieParam, VIDEO_FRAME_INFO_S *pstFrmInfo, HI_U32 u32SegIdx, HI_U32 u32NodeIdx);

HI_S32 SVP_NNIE_Yolo_Init(SVP_NNIE_PARAM_S *pstNnieParam, SVP_YOLO_PARAM_S *pstYoloParam);

HI_VOID SVP_NNIE_Yolo_DeInit(SVP_YOLO_PARAM_S *pstYoloParam);

HI_S32 SVP_YOLO_PostProcess(SVP_NNIE_PARAM_S *pstNnieParam, SVP_YOLO_PARAM_S *pstYoloParam, HI_BOOL bVerbose);

#endif